// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: hyuan

#include "coin2/exchange/bybit/api_util/rate_tracker.h"

namespace coin2::exchange::bybit::api_util {

constexpr auto EXAMPLE_SPOT_BYBIT = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke18/bybit/trade_key.json",
  "rate_limit_info": {
    "spot_bybit": {
      "get_method": {
        "secs": 1,
        "reqs": 70
      },
      "post_method": {
        "secs": 1,
        "reqs": 50
      },
      "endpoint": {
        "secs": 1,
        "reqs": 20
      }
    }
  }
},)";

template <class RateTracker>
void CheckFail(RateTracker&& rate_tracker, bool check_result) {
  if (!check_result) {
    LOG(INFO) << "Hit rate limit on [" << rate_tracker->GetName() << "]";
  }
}

BybitRateTracker::BybitRateTracker(const MarketExchangeApi& mea,
    const coin2::exchange::order::OrderConnectionConfig& cfg) {
  if (mea.market == coin::proto::MarketType::Spot) {
    RateLimitInfoSpotBybit rate_config;
    RateConfig::Get()->GetConfig(mea, cfg, &rate_config);
    InitRateTracker(&rate_config, mea.String(), EXAMPLE_SPOT_BYBIT);
  } else {
    CHECK(false) << "unsupported mea: " << mea.String();
  }
}

void BybitRateTracker::InitRateTracker(
    RateLimitInfoSpotBybit* rate_config,
    std::string_view mea_str,
    std::string_view example) {
  if (!rate_config || !rate_config->has_get_method() ||
      !rate_config->has_post_method() ||
      !rate_config->has_endpoint()) {
    coin2::exchange::base::order::RateConfig::CheckFail(mea_str, example);
  }

  const auto& get_method = rate_config->get_method();
  const auto& post_method = rate_config->post_method();
  const auto& endpoint = rate_config->endpoint();
  LOG(INFO) << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
      << "\nOG request rate was set as below:\n"
      << "\nget_method: [" << get_method.reqs() << "(reqs)/" << get_method.secs() << "(secs)]"
      << "\npost_method: [" << post_method.reqs() << "(reqs)/" << post_method.secs() << "(secs)]"
      << "\nendpoint: [" << endpoint.reqs() << "(reqs)/" << endpoint.secs() << "(secs)]"
      << "\n\nOr you can configure customized rate in driver config. For example: " << example
      << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~";

  rate_tracker_getmethod_ =
      std::make_unique<FixedWindowRateTracker>(get_method.secs() * 1'000'000'000LL);
  rate_tracker_getmethod_->SetRateLimitInfo(get_method.reqs(), get_method.reqs());
  rate_tracker_getmethod_->SetName("bybit_rate_tracker_getmethod(1 sec)");

  rate_tracker_postmethod_ =
      std::make_unique<FixedWindowRateTracker>(post_method.secs() * 1'000'000'000LL);
  rate_tracker_postmethod_->SetRateLimitInfo(post_method.reqs(), post_method.reqs());
  rate_tracker_postmethod_->SetName("bybit_rate_tracker_postmethod(1 sec)");

  rate_tracker_balance_ =
      std::make_unique<FixedWindowRateTracker>(endpoint.secs() * 1'000'000'000LL);
  rate_tracker_balance_->SetRateLimitInfo(endpoint.reqs(), endpoint.reqs());
  rate_tracker_balance_->SetName("bybit_rate_tracker_balance(1 sec)");

  rate_tracker_openorder_ =
      std::make_unique<FixedWindowRateTracker>(endpoint.secs() * 1'000'000'000LL);
  rate_tracker_openorder_->SetRateLimitInfo(endpoint.reqs(), endpoint.reqs());
  rate_tracker_openorder_->SetName("bybit_rate_tracker_openorder(1 sec)");

  rate_tracker_historder_ =
      std::make_unique<FixedWindowRateTracker>(endpoint.secs() * 1'000'000'000LL);
  rate_tracker_historder_->SetRateLimitInfo(endpoint.reqs(), endpoint.reqs());
  rate_tracker_historder_->SetName("bybit_rate_tracker_historder(1 sec)");

  rate_tracker_fillorder_ =
      std::make_unique<FixedWindowRateTracker>(endpoint.secs() * 1'000'000'000LL);
  rate_tracker_fillorder_->SetRateLimitInfo(endpoint.reqs(), endpoint.reqs());
  rate_tracker_fillorder_->SetName("bybit_rate_tracker_fillorder(1 sec)");

  rate_tracker_placeorder_ =
      std::make_unique<FixedWindowRateTracker>(endpoint.secs() * 1'000'000'000LL);
  rate_tracker_placeorder_->SetRateLimitInfo(endpoint.reqs(), endpoint.reqs());
  rate_tracker_placeorder_->SetName("bybit_rate_tracker_placeorder(1 sec)");

  rate_tracker_cancelorder_ =
      std::make_unique<FixedWindowRateTracker>(endpoint.secs() * 1'000'000'000LL);
  rate_tracker_cancelorder_->SetRateLimitInfo(endpoint.reqs(), endpoint.reqs());
  rate_tracker_cancelorder_->SetName("bybit_rate_tracker_cancelorder(1 sec)");

  rate_limit_report_writer_ = std::make_unique<RateLimitReportWriter>();
  rate_limit_report_writer_->SetRateLimitInfoString(rate_config);
}

bool BybitRateTracker::HasEnoughRateLimit(int64_t timestamp, ActionType action_type,
    QueryType query_type) {
  bool ret = true;
  if (action_type == ActionType::PLACE) {
    if (rate_tracker_postmethod_) {
      ret = rate_tracker_postmethod_->GetRemaining(timestamp) > 2;
      CheckFail(rate_tracker_postmethod_, ret);
    }
    if (rate_tracker_placeorder_) {
      ret = rate_tracker_placeorder_->GetRemaining(timestamp) > 1;
      CheckFail(rate_tracker_placeorder_, ret);
    }
  } else if (action_type == ActionType::CANCEL) {
    if (rate_tracker_cancelorder_) {
      ret = rate_tracker_cancelorder_->GetRemaining(timestamp) > 1;
      CheckFail(rate_tracker_cancelorder_, ret);
    }
  } else if (action_type == ActionType::QUERY) {
    if (rate_tracker_getmethod_) {
      ret = rate_tracker_getmethod_->GetRemaining(timestamp) > 2;
      CheckFail(rate_tracker_getmethod_, ret);
    }
    if (query_type == QueryType::ACCOUNT_BALANCE) {
      if (rate_tracker_balance_) {
        ret = rate_tracker_balance_->GetRemaining(timestamp) > 1;
        CheckFail(rate_tracker_balance_, ret);
      }
    } else if (query_type == QueryType::OPEN_ORDER) {
      if (rate_tracker_openorder_) {
        ret = (rate_tracker_openorder_->GetRemaining(timestamp) > 2);
        CheckFail(rate_tracker_openorder_, ret);
      }
    } else if (query_type == QueryType::HIST_ORDER) {
      if (rate_tracker_historder_) {
        ret = (rate_tracker_historder_->GetRemaining(timestamp) > 2);
        CheckFail(rate_tracker_historder_, ret);
      }
    } else if (query_type == QueryType::FILL_ORDER) {
      if (rate_tracker_fillorder_) {
        ret = (rate_tracker_fillorder_->GetRemaining(timestamp) > 2);
        CheckFail(rate_tracker_fillorder_, ret);
      }
    }
  }

  if (!ret) {
    rate_limit_report_writer_->SetHitRateLimit(timestamp, action_type);
  }
  return ret;
}

void BybitRateTracker::Push(int32_t num_reqs, int64_t timestamp, ActionType action_type,
    QueryType query_type) {
  if (action_type == ActionType::PLACE) {
    if (rate_tracker_postmethod_) {
      rate_tracker_postmethod_->Push(num_reqs, timestamp);
    }
    if (rate_tracker_placeorder_) {
      rate_tracker_placeorder_->Push(num_reqs, timestamp);
    }
  } else if (action_type == ActionType::CANCEL) {
    if (rate_tracker_cancelorder_) {
      rate_tracker_cancelorder_->Push(num_reqs, timestamp);
    }
  } else if (action_type == ActionType::QUERY) {
    if (rate_tracker_getmethod_) {
      rate_tracker_getmethod_->Push(num_reqs, timestamp);
    }
    if (query_type == QueryType::ACCOUNT_BALANCE) {
      if (rate_tracker_balance_) {
        rate_tracker_balance_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::OPEN_ORDER) {
      if (rate_tracker_openorder_) {
        rate_tracker_openorder_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::HIST_ORDER) {
      if (rate_tracker_historder_) {
        rate_tracker_historder_->Push(num_reqs, timestamp);
      }
    } else if (query_type == QueryType::FILL_ORDER) {
      if (rate_tracker_fillorder_) {
        rate_tracker_fillorder_->Push(num_reqs, timestamp);
      }
    } else {
      return;
    }
  } else {
    return;
  }
}

void BybitRateTracker::SetRateLimitInfoString(const google::protobuf::Message* message) {
  if (!rate_limit_report_writer_) {
    return;
  }
  rate_limit_report_writer_->SetRateLimitInfoString(message);
}

void BybitRateTracker::WriteRateLimitReportMessage(OrderLoggerCB callback) {
  if (!callback) {
    return;
  }
  rate_limit_report_writer_->Write(callback);
}

double BybitRateTracker::GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) {
  switch (action_type) {
    case ActionType::PLACE: {
      return GetMaxRateUsedRatio(rate_tracker_placeorder_, timestamp);
    }
    case ActionType::CANCEL: {
      return GetMaxRateUsedRatio(rate_tracker_cancelorder_, timestamp);
    }
    case ActionType::QUERY: {
      double balance_rate_used_ratios =
          GetMaxRateUsedRatio(rate_tracker_balance_, timestamp);
      double openorder_rate_used_ratios =
          GetMaxRateUsedRatio(rate_tracker_openorder_, timestamp);
      double historder_rate_used_ratios =
          GetMaxRateUsedRatio(rate_tracker_historder_, timestamp);
      double fillorder_rate_used_ratios =
          GetMaxRateUsedRatio(rate_tracker_fillorder_, timestamp);
      return std::max(std::max(balance_rate_used_ratios, openorder_rate_used_ratios),
          std::max(historder_rate_used_ratios, fillorder_rate_used_ratios));
    }
    default: {
      // shouldn't reached here
      LOG(ERROR) << "ops~ something is wrong, action_type=" << action_type;
      return 0.;
    }
  }
}

}  // namespace coin2::exchange::bybit::api_util
